# -- coding: utf-8 --
"""
alpha方程求解
"""

# third party lib
import numpy as np

# private lib
from utils.constant import InterpMode
from utils.numeric import grid_extrap, penta_diagonal_solve, grid_interp


def cal_alpha_g(params: dict, liquid_vapour_mass_transfer_rate: np.ndarray) -> np.ndarray:
    """连续性方程, 计算气相截面分数, 二阶迎风格式
    Args:
        params: dict 计算需要的输入参数
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率
        
    Returns:
        alpha_g_node_LHS: numpy ndarray float64, 气相截面分数
    """
    # 直接从params中提取所需变量
    dx = params["dx"]
    dt = params["dt"]
    V_g_segment = params["V_g_segment"]
    rho_g_node = params["rho_g_node"]
    rho_g_node_iteration_init = params["rho_g_node_iteration_init"]
    alpha_g_node_iteration_init = params["alpha_g_node_iteration_init"]
    alpha_g_node = params["alpha_g_node"]
    
    # 速度插值
    V_g_node = grid_extrap(V_g_segment, 1)
    
    # 相变项
    psi_g_node = grid_extrap(liquid_vapour_mass_transfer_rate, 1)
    
    # 计算正负速度 
    V_g_node_plus = np.maximum(V_g_node[1:-1], 0)
    V_g_node_neg = - np.maximum(-V_g_node[1:-1], 0)
    
    # 计算内部点系数（二阶迎风）
    a_P = dx/dt - 3*V_g_node_neg/2 + 3*V_g_node_plus/2 + (V_g_segment[1:] - V_g_segment[:-1])
    a_E = 4 * V_g_node_neg/2
    a_EE = - V_g_node_neg/2
    a_W = - 4 * V_g_node_plus/2
    a_WW = V_g_node_plus/2
    b = (dx * psi_g_node + dx/dt*(alpha_g_node_iteration_init * rho_g_node_iteration_init))[1:-1]
    
    # 计算边界系数（一阶迎风）
    a_P[0] = dx/dt + V_g_node_plus[0] - V_g_node_neg[0] + V_g_segment[1] - V_g_segment[0]
    a_P[-1] = dx/dt + V_g_node_plus[-1] - V_g_node_neg[-1] + V_g_segment[-1] - V_g_segment[-2]
    
    a_W[0] = - V_g_node_plus[0]
    a_W[-1] = - V_g_node_plus[-1]
    
    a_E[0] = V_g_node_neg[0]
    a_E[-1] = V_g_node_neg[-1]
    
    a_WW[0] = 0
    a_WW[-1] = 0
    
    a_EE[0] = 0
    a_EE[-1] = 0

    # 增加边界效应
    alpha_rho_g_node = alpha_g_node * rho_g_node
    b[0] -= a_W[0]*alpha_rho_g_node[0]
    b[-1] -= a_E[-1]*alpha_rho_g_node[-1]
    b[1] -= a_WW[1]*alpha_rho_g_node[0]
    b[-2] -= a_EE[-2]*alpha_rho_g_node[-1]
    
    # 求解五对角矩阵
    alpha_rho_g_node[1:-1] = penta_diagonal_solve(a_P, a_E, a_W, a_EE, a_WW, b)
    alpha_g_node_update = alpha_rho_g_node/rho_g_node
    return np.clip(alpha_g_node_update, 0, 1)   

def cal_alpha_L(params: dict, liquid_vapour_mass_transfer_rate: np.ndarray, liquid_to_droplet_mass_transfer_rate: np.ndarray, droplet_to_liquid_mass_transfer_rate: np.ndarray) -> np.ndarray:
    """连续性方程，计算截面分数，液相
    Args:
        params: dict 计算需要的输入参数
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率
        liquid_to_droplet_mass_transfer_rate: np.ndarray 液相到液滴质量传递速率
        droplet_to_liquid_mass_transfer_rate: np.ndarray 液滴到液相质量传递速率
    Returns:
        alpha_L_node_LHS: numpy ndarray float64, 液相截面分数
    """
    # 直接从params中提取所需变量
    dx = params["dx"]
    dt = params["dt"]

    V_L_segment = params["V_L_segment"]
    rho_L_node = params["rho_L_node"]
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]
    alpha_L_node_iteration_init = params["alpha_L_node_iteration_init"]
    alpha_L_node = params["alpha_L_node"]
    alpha_D_node = params["alpha_D_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)

    # 速度插值
    V_L_node = grid_extrap(V_L_segment, 1)
    
    # 计算相变项
    psi_L_node = grid_extrap(- liquid_vapour_mass_transfer_rate * (alpha_L_segment / (alpha_L_segment + alpha_D_segment)) - liquid_to_droplet_mass_transfer_rate + droplet_to_liquid_mass_transfer_rate, 1)
    
    # 计算正负速度
    V_L_node_plus = np.maximum(V_L_node[1:-1], 0)
    V_L_node_neg = - np.maximum(- V_L_node[1:-1], 0)
    
    # 计算内部点系数（二阶迎风）
    a_P = - 3*V_L_node_neg/2 + 3*V_L_node_plus/2 + dx/dt + (V_L_segment[1:] - V_L_segment[:-1])
    a_E = 4 * V_L_node_neg/2
    a_EE = - V_L_node_neg/2
    a_W = - 4 * V_L_node_plus/2
    a_WW = V_L_node_plus/2
    b = (dx * psi_L_node + dx/dt*(alpha_L_node_iteration_init * rho_L_node_iteration_init))[1:-1]
    
    # 计算边界系数（一阶迎风）
    a_P[0] = V_L_node_plus[0] - V_L_node_neg[0] + (dx/dt + V_L_segment[1] - V_L_segment[0])
    a_P[-1] = V_L_node_plus[-1] - V_L_node_neg[-1] + (dx/dt + V_L_segment[-1] - V_L_segment[-2])
    
    a_W[0] = - V_L_node_plus[0]
    a_W[-1] = - V_L_node_plus[-1]
    
    a_E[0] = V_L_node_neg[0]
    a_E[-1] = V_L_node_neg[-1]
    
    a_WW[0] = 0
    a_WW[-1] = 0
    
    a_EE[0] = 0
    a_EE[-1] = 0

    # 增加边界效应
    alpha_rho_L_node = np.copy(alpha_L_node * rho_L_node)
    b[0] -= a_W[0] * alpha_rho_L_node[0]
    b[-1] -= a_E[-1] * alpha_rho_L_node[-1]
    b[1] -= a_WW[1] * alpha_rho_L_node[0]
    b[-2] -= a_EE[-2] * alpha_rho_L_node[-1]
    
    # 求解五对角矩阵
    alpha_rho_L_node[1:-1] = penta_diagonal_solve(a_P, a_E, a_W, a_EE, a_WW, b)
    alpha_L_node_update = alpha_rho_L_node/rho_L_node
    return np.clip(alpha_L_node_update, 0, 1)   

def cal_alpha_D(params: dict, liquid_vapour_mass_transfer_rate: np.ndarray, liquid_to_droplet_mass_transfer_rate: np.ndarray, droplet_to_liquid_mass_transfer_rate: np.ndarray, droplet_velocity: np.ndarray) -> np.ndarray:
    """连续性方程，计算截面分数，液滴相
    Args:
        params: dict 计算需要的输入参数
        liquid_vapour_mass_transfer_rate: np.ndarray 液相蒸气质量传递速率
        liquid_to_droplet_mass_transfer_rate: np.ndarray 液相到液滴质量传递速率
        droplet_to_liquid_mass_transfer_rate: np.ndarray 液滴到液相质量传递速率
        droplet_velocity: np.ndarray 液滴速度
    Returns:
        alpha_D_node_LHS: numpy ndarray float64, 液滴截面分数
    """
    # 直接从params中提取所需变量
    dt = params["dt"]
    dx = params["dx"]


    rho_L_node = params["rho_L_node"]
    rho_L_node_iteration_init = params["rho_L_node_iteration_init"]
    alpha_D_node_iteration_init = params["alpha_D_node_iteration_init"]
    alpha_L_node = params["alpha_L_node"]
    alpha_D_node = params["alpha_D_node"]
    alpha_L_segment = grid_interp(alpha_L_node, InterpMode.MID)
    alpha_D_segment = grid_interp(alpha_D_node, InterpMode.MID)
    
    # 相变项计算
    mass_src_D = (
        -liquid_vapour_mass_transfer_rate * (alpha_D_segment / (alpha_L_segment + alpha_D_segment))
        + liquid_to_droplet_mass_transfer_rate
        - droplet_to_liquid_mass_transfer_rate
    )
    mass_src_D_node = grid_extrap(mass_src_D, 1)
    
    # 速度插值
    V_D_segment = np.copy(droplet_velocity)
    V_D_node = grid_extrap(V_D_segment, 1)
    
    # 计算正负速度
    V_D_node_plus = np.maximum(V_D_node[1:-1], 0)
    V_D_node_neg = - np.maximum(- V_D_node[1:-1], 0)
    
    # 计算内部点系数（二阶迎风）
    a_P = - 3*V_D_node_neg/2 + 3*V_D_node_plus/2 + dx/dt + (V_D_segment[1:] - V_D_segment[:-1])
    a_E = 4 * V_D_node_neg/2
    a_EE = - V_D_node_neg/2
    a_W = - 4 * V_D_node_plus/2
    a_WW = V_D_node_plus/2
    b = (dx * mass_src_D_node + dx/dt*(alpha_D_node_iteration_init * rho_L_node_iteration_init))[1:-1]
    
    # 计算边界系数（一阶迎风）
    a_P[0] = V_D_node_plus[0] - V_D_node_neg[0] + (dx/dt + V_D_segment[1] - V_D_segment[0])
    a_P[-1] = V_D_node_plus[-1] - V_D_node_neg[-1] + (dx/dt + V_D_segment[-1] - V_D_segment[-2])
    
    a_W[0] = - V_D_node_plus[0]
    a_W[-1] = - V_D_node_plus[-1]
    
    a_E[0] = V_D_node_neg[0]
    a_E[-1] = V_D_node_neg[-1]
    
    a_WW[0] = 0
    a_WW[-1] = 0
    
    a_EE[0] = 0
    a_EE[-1] = 0

    # 增加边界效应
    alpha_rho_D_node = np.copy(alpha_D_node * rho_L_node)
    b[0] -= a_W[0] * alpha_rho_D_node[0]
    b[-1] -= a_E[-1] * alpha_rho_D_node[-1]
    b[1] -= a_WW[1] * alpha_rho_D_node[0]
    b[-2] -= a_EE[-2] * alpha_rho_D_node[-1]
    
    # 求解五对角矩阵
    alpha_rho_D_node[1:-1] = penta_diagonal_solve(a_P, a_E, a_W, a_EE, a_WW, b)
    alpha_D_node_update = alpha_rho_D_node/rho_L_node
    return np.clip(alpha_D_node_update, 0, 1)   
